.globl classify

.text
classify:
    # =====================================
    # COMMAND LINE ARGUMENTS
    # =====================================
    # Args:
    #   a0 (int)    argc
    #   a1 (char**) argv
    #   a2 (int)    print_classification, if this is zero,
    #               you should print the classification. Otherwise,
    #               this function should not print ANYTHING.
    # Returns:
    #   a0 (int)    Classification
    #
    # If there are an incorrect number of command line args,
    # this function returns with exit code 49.
    #
    # Usage:
    #   main.s -m -1 <M0_PATH> <M1_PATH> <INPUT_PATH> <OUTPUT_PATH>

    # prologue
    addi sp, sp, -32
    sw ra, 0(sp)
    sw s0, 4(sp)
    sw s1, 8(sp)
    sw s2, 12(sp)
    sw s3, 16(sp)
    sw s4, 20(sp)
    sw s5, 24(sp)
    sw s6, 28(sp)

    # check off the argc
    # if a0 == 5, goto load_mat
    addi t0, x0, 5
    beq a0, t0, load_mat
    addi a1, x0, 49
    jal ra, exit2


    # =====================================
    # LOAD MATRICES
    # =====================================
load_mat:
    # save arguments
    lw s0, 4(a1)   # pointer to <M0_PATH>
    lw s1, 8(a1)   # pointer to <M1_PATH>
    lw s2, 12(a1)  # pointer to <INPUT_PATH>
    lw s3, 16(a1)  # pointer to <OUTPUT_PATH>
    add s4, a2, x0 # if to print classification

    # Load pretrained m0
    addi sp, sp, -8
    add a0, s0, x0
    addi a1, sp, 0 # 16(sp) = m0_row
    addi a2, sp, 4 # 20(sp) = m0_col
    jal ra, read_matrix
    add s0, a0, x0 # s0 = pointer to m0

    # Load pretrained m1
    addi sp, sp, -8
    add a0, s1, x0
    addi a1, sp, 0 # 8(sp)  = m1_row
    addi a2, sp, 4 # 12(sp) = m1_col
    jal ra, read_matrix
    add s1, a0, x0 # s1 = pointer to m1

    # Load input matrix
    addi sp, sp, -8
    add a0, s2, x0
    addi a1, sp, 0 # 0(sp) = input_row
    addi a2, sp, 4 # 4(sp) = input_col
    jal ra, read_matrix
    add s2, a0, x0 # s2 = pointer to input


    # =====================================
    # RUN LAYERS
    # =====================================
    # 1. LINEAR LAYER:    m0 * input
    # 2. NONLINEAR LAYER: ReLU(m0 * input)
    # 3. LINEAR LAYER:    m1 * ReLU(m0 * input)

    # allocate rooms for m0 * input result
    lw t0, 16(sp)
    lw t1, 4(sp)
    mul a0, t0, t1
    slli a0, a0, 2
    jal ra, malloc
    add s5, a0, x0

    # m0 * input
    add a0, s0, x0
    lw a1, 16(sp)
    lw a2, 20(sp)
    add a3, s2, x0
    lw a4, 0(sp)
    lw a5, 4(sp)
    add a6, s5, x0 # s5 pointer to m0 * input
    jal ra, matmul

    # ReLU(m0 * input)
    add a0, s5, x0
    lw t0, 16(sp)
    lw t1, 4(sp)
    mul a1, t0, t1
    jal ra, relu

    # allocate rooms for m1 * ReLU result
    lw t0, 8(sp)
    lw t1, 4(sp)
    mul a0, t0, t1
    slli a0, a0, 2
    jal ra, malloc
    add s6, a0, x0

    # m1 * ReLU
    add a0, s1, x0
    lw a1, 8(sp)
    lw a2, 12(sp)
    add a3, s5, x0
    lw a4, 16(sp)
    lw a5, 4(sp)
    add a6, s6, x0 # s6 = m1 * ReLU
    jal ra, matmul

    # free s5(m0 * input)
    add a0, s5, x0
    jal ra, free

    # =====================================
    # WRITE OUTPUT
    # =====================================
    # Write output matrix
    add a0, s3, x0
    add a1, s6, x0
    lw a2, 8(sp)
    lw a3, 4(sp)
    jal ra, write_matrix

    # =====================================
    # CALCULATE CLASSIFICATION/LABEL
    # =====================================
    # Call argmax
    add a0, s6, x0
    lw t0, 8(sp)
    lw t1, 4(sp)
    mul a1, t0, t1
    jal ra, argmax
    # save classification
    add s5, a0, x0

    # Print classification
    # if s4 != 0, goto end
    bne s4, x0, end
    add a1, s5, x0
    jal ra, print_int

    # Print newline afterwards for clarity
    addi a1, x0, '\n'
    jal ra, print_char

end:
    addi sp, sp, 24
    # free s6(m1 * ReLU)
    add a0, s6, x0
    jal ra, free
    # set return value
    add a0, s5, x0

    # epilogue
    lw ra, 0(sp)
    lw s0, 4(sp)
    lw s1, 8(sp)
    lw s2, 12(sp)
    lw s3, 16(sp)
    lw s4, 20(sp)
    lw s5, 24(sp)
    lw s6, 28(sp)
    addi sp, sp, 32

    ret
